"""
恶意代码检测系统 - 数据预处理模块
功能：实现数据加载、清洗、标准化等预处理流程
"""

import pandas as pd
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import train_test_split
from data.collector.virustotal_downloader import VirusTotalCollector
from data.scripts.cleaner import DataCleaner
from data.scripts.labeler import DataLabeler
from data.scripts.balancer import DataBalancer

def load_raw_data(file_path):
    """加载原始恶意软件特征数据
    参数：
        file_path: 原始数据文件路径（CSV格式）
    返回：
        pandas.DataFrame 包含原始特征数据
    """
    return pd.read_csv(file_path)

def clean_data(df):
    """数据清洗流程
    1. 缺失值填充为0
    2. 去除重复样本
    参数：
        df: 原始数据DataFrame
    返回：
        清洗后的DataFrame
    """
    df.fillna(0, inplace=True)  # 用0填充缺失值
    df.drop_duplicates(inplace=True)  # 删除重复行
    return df

def normalize_features(df):
    """特征标准化处理
    使用Z-score标准化方法对特征进行归一化
    参数：
        df: 清洗后的DataFrame
    返回：
        标准化后的特征DataFrame（不包含标签列）
    """
    scaler = StandardScaler()
    scaled_features = scaler.fit_transform(df.iloc[:, :-1])  # 最后一列为标签
    return pd.DataFrame(scaled_features, columns=df.columns[:-1])

def preprocess_pipeline(input_path, output_path):
    """完整预处理流水线
    参数：
        input_path: 原始数据输入路径
        output_path: 处理后的数据保存路径
    """
    raw_df = load_raw_data(input_path)
    cleaned_df = clean_data(raw_df)
    processed_df = normalize_features(cleaned_df)
    processed_df['label'] = cleaned_df['label']  # 重新添加标签列
    processed_df.to_csv(output_path, index=False)

def full_pipeline():
    """端到端数据处理流程
    包含：数据收集->清洗->标注->平衡
    """
    # 1. 从VirusTotal收集数据
    vt = VirusTotalCollector(API_KEY)
    vt.download_samples("tag:ransomware", "data/raw/ransomware")
    
    # 2. 数据清洗
    cleaner = DataCleaner()
    cleaner.clean_dataset("data/raw", "data/clean")
    
    # 3. 自动标注
    labeler = DataLabeler("threat_intelligence.csv")
    labeler.label_data("data/clean", "data/labeled.csv")
    
    # 4. 处理数据不平衡
    df = pd.read_csv("data/labeled.csv")
    X, y = df.drop('is_malicious', axis=1), df['is_malicious']
    balancer = DataBalancer(strategy='smote')
    X_balanced, y_balanced = balancer.balance(X, y) 